#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/3/6 16:07
# @Author  : wanghaoran
# @File    : qwen_72b.py
import requests
import json


def chat_qwen_72b(prompt, if_print=False):
    """
    :param database_schema: 数据库结构
    :param question: 问题
    :param if_print: 是否打印
    :return:
    """
    url = "http://180.184.40.76:11119/v1/chat/completions"
    payload = {
        "model": "Qwen/Qwen-72B-Chat",
        "temperature": 0,
        "do_sample": True,
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ],
        "stream": True,
        "max_tokens": 1024,
        "stop": ["<|endoftext|>", "<|im_start|>", "<|im_end|>"]
    }
    headers = {
        "Content-Type": "application/json"
    }
    response = requests.request(
        "POST",
        url,
        headers=headers,
        json=payload,
        stream=True,
        timeout=20
        )
    whole_text = ""
    token_num = 0
    for chunk in response.iter_lines(chunk_size=1024):
        chunk = chunk.decode('UTF-8')
        if chunk and "data: [DONE]" not in chunk:
            delta = json.loads(chunk.replace("data: ", ""))
            delta = delta["choices"][0]["delta"]
            if "content" in delta:
                cur_token = delta["content"]
                token_num += 1
                if if_print:
                    print(cur_token, end="", flush=True)
                whole_text += cur_token
    return whole_text
